import math
from collections import OrderedDict
import torch.nn as nn
import numpy as np
import cv2


def showcharacterimage(tensor):
    image = tensor.cpu()
    image = image.detach().numpy()
    image = np.clip(image, 0, 100)
    image = 255 * ((image - image.min()) / (image.max() - image.min()))
    image = image[0, 0, :, :]
    cv2.imshow("", image.astype(np.uint8))

# ---------------------------------------------------------------------#
#   残差结构
#   利用一个1x1卷积下降通道数，然后利用一个3x3卷积提取特征并且上升通道数
#   最后接上一个残差边
# ---------------------------------------------------------------------#
# class MCBLBlock(nn.Module):
#     def __init__(self, in_channels, out_channels):
#         super(MCBLBlock, self).__init__()
#         self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
#         self.depth_conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1, bias=False,
#                                     groups=in_channels)
#         self.point_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=False)
#         self.bn = nn.BatchNorm2d(num_features=out_channels)
#         self.activate = nn.LeakyReLU(negative_slope=0.1)
#
#     def forward(self, x):
#         x = self.pool(x)
#         x = self.depth_conv(x)
#         out = self.point_conv(x)
#         x = self.bn(out)
#         x = self.activate(x)
#         return x, out

class MCBLBlock(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(MCBLBlock, self).__init__()
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=False)
        self.bn = nn.BatchNorm2d(num_features=out_channels)
        self.activate = nn.LeakyReLU(negative_slope=0.1)

    def forward(self, x):
        x = self.pool(x)
        out = self.conv(x)
        x = self.bn(out)
        x = self.activate(x)
        return x, out

# class CBLBlock(nn.Module):
#     def __init__(self, in_channels, out_channels, kernel_size, padding):
#         super(CBLBlock, self).__init__()
#         self.depth_conv = nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=False,
#                                     groups=in_channels)
#         self.point_conv = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=False)
#         self.bn = nn.BatchNorm2d(num_features=out_channels)
#         self.activate = nn.LeakyReLU(negative_slope=0.1)
#
#     def forward(self, x):
#         x = self.depth_conv(x)
#         x = self.point_conv(x)
#         x = self.bn(x)
#         x = self.activate(x)
#         return x

class CBLBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding):
        super(CBLBlock, self).__init__()
        self.conv = nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=kernel_size, stride=1, padding=padding, bias=False)
        self.bn = nn.BatchNorm2d(num_features=out_channels)
        self.activate = nn.LeakyReLU(negative_slope=0.1)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.activate(x)
        return x

# noinspection PyListCreation
class DarkNet(nn.Module):
    def __init__(self):
        super(DarkNet, self).__init__()

        self.conv1 = CBLBlock(in_channels=3, out_channels=16, kernel_size=3, padding=1)

        self.mcbl1 = MCBLBlock(in_channels=16, out_channels=32)
        self.mcbl2 = MCBLBlock(in_channels=32, out_channels=64)
        self.mcbl3 = MCBLBlock(in_channels=64, out_channels=128)
        self.mcbl4 = MCBLBlock(in_channels=128, out_channels=256)
        self.mcbl5 = MCBLBlock(in_channels=256, out_channels=512)

        # padding右边和下边
        self.paddrd = nn.ZeroPad2d(padding=(0, 1, 0, 1))
        self.pool = nn.MaxPool2d(kernel_size=2, stride=1)

        self.cbl1 = CBLBlock(in_channels=512, out_channels=1024, kernel_size=3, padding=1)

        # 进行权值初始化
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
                m.weight.data.normal_(0, math.sqrt(2. / n))
            elif isinstance(m, nn.BatchNorm2d):
                m.weight.data.fill_(1)
                m.bias.data.zero_()

    def forward(self, x):

        x = self.conv1(x)
        # showcharacterimage(x)
        x, _ = self.mcbl1(x)
        x, _ = self.mcbl2(x)
        x, _ = self.mcbl3(x)
        x, out2 = self.mcbl4(x)

        x, _ = self.mcbl5(x)
        x = self.paddrd(x)
        x = self.pool(x)
        out1 = self.cbl1(x)


        return out1, out2

def darknet53():
    model = DarkNet()

    return model
